#include <Session.h>
#include <utils.h>

#ifdef _WIN32
#pragma comment(lib, "ws2_32.lib")
#define INVALID_SOCKET_HANDLE INVALID_SOCKET
#define SOCKET_ERROR_CODE WSAGetLastError()
#else
#define INVALID_SOCKET_HANDLE -1
#define SOCKET_ERROR_CODE errno
#endif
#define USER_PROMPT "->"
constexpr static unsigned int SESSION_RECV_BUFFER_SIZE = 4096;

Session::Session(SocketHandle sock, const struct sockaddr_in& addr)
{
    m_socket = sock;
    m_address = addr;
    m_readBuffer.clear();
    m_writeBuffer.clear();
    m_isClosed = false;
    m_authStage = AUTH_INIT;
    m_username.clear();
    m_permission = UserPermission::BANNED;
    m_prompt = USER_PROMPT;
    m_loginRetryCount = 0;
    resetHeartbeatTimer(SESSION_LOGIN_TIMEOUT);
}

Session::~Session() 
{
    close();
}
AuthStage Session::getAuthStage() const
{
    return m_authStage;
}
void Session::setAuthStage(AuthStage stage)
{
    m_authStage = stage;
}
bool Session::isAuthenticated() const
{
    return (m_authStage & AuthStage::AUTH_SUCCESS) == AuthStage::AUTH_SUCCESS;
}
const std::string &Session::getUsername() const
{
    return m_username;
}
void Session::setUsername(const std::string &username)
{
    m_username = username;
    m_prompt = utils::format("%s%s", username.c_str(), USER_PROMPT);
}
const UserPermission &Session::getPermission() const
{
    return m_permission;
}
void Session::setPermission(UserPermission permission)
{
    m_permission = permission;
}
void Session::resetHeartbeatTimer(int timeSec)
{
    m_heartbeatTimeout = timeSec;
}
bool Session::checkHeartbeatValidity() const
{
    return m_heartbeatTimeout > 0;
}
void Session::refreshHeartbeatTimer()
{
    if (m_heartbeatTimeout > 0)
        m_heartbeatTimeout--;
}
void Session::refreshAuthFail()
{
    m_loginRetryCount++;
    if (m_loginRetryCount >= 3)
    {
        appendWriteData("\nLogin failed (3 attempts exceeded). Connection closed.\n");
        send();
        close();
    }
    else
    {
        appendWriteData("\nInvalid username or password! Remaining attempts: " +
            std::to_string(3 - m_loginRetryCount) + "\n");
        setAuthStage(AuthStage::AUTH_WAIT_USER);
    }
}
SocketHandle Session::getSocket() const 
{
    return m_socket;
}

const struct sockaddr_in& Session::getAddress() const 
{
    return m_address;
}

std::string Session::getAddressStr() const 
{
    char ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(m_address.sin_addr), ip, INET_ADDRSTRLEN);
    return std::string(ip) + ":" + std::to_string(ntohs(m_address.sin_port));
}

const std::string Session::getUsrPrompt() const
{
    return m_prompt;
}

void Session::extractRowData(std::string& cmdStr)
{
    std::lock_guard<std::mutex> lock(m_readLock);
    cmdStr = m_readBuffer;
    m_readBuffer.clear();
}

void Session::appendReadData(const std::string& data)
{
    appendReadData(data.data(), data.size());
}

void Session::appendWriteData(const std::string& data)
{
    appendWriteData(data.data(), data.size());
}

void Session::appendReadData(const char* data, size_t len) 
{
    std::lock_guard<std::mutex> lock(m_readLock);
    if (!isClosed()) m_readBuffer.append(data, len);
}

void Session::appendWriteData(const char* data, size_t len)
{
    std::lock_guard<std::mutex> lock(m_writeLock);
    if (!isClosed()) m_writeBuffer.append(data, len);
}

void Session::clearReadBuffer() 
{
    std::lock_guard<std::mutex> lock(m_readLock);
    m_readBuffer.clear();
}

void Session::clearWriteBuffer()
{
    std::lock_guard<std::mutex> lock(m_writeLock);
    m_writeBuffer.clear();
}

bool Session::isClosed() const 
{
    return m_isClosed;
}

bool Session::isWritable() const
{
    std::lock_guard<std::mutex> lock(m_writeLock);

    return m_writeBuffer.empty();
}

bool Session::isReadable() const
{
    std::lock_guard<std::mutex> lock(m_readLock);

    return m_readBuffer.empty();
}

void Session::close() 
{
    if (!m_isClosed && m_socket != INVALID_SOCKET_HANDLE) 
    {
        CLOSE_SOCKET(m_socket);
        m_socket = INVALID_SOCKET_HANDLE;
        m_isClosed = true;
    }
}

ssize_t Session::recv()
{
    char buffer[SESSION_RECV_BUFFER_SIZE] = { 0 };

    ssize_t bytesRead = ::recv(m_socket, buffer, SESSION_RECV_BUFFER_SIZE - 1, 0);

    if (bytesRead > 0)
    {
        appendReadData(buffer, bytesRead);
    }

    return bytesRead;
}

ssize_t Session::send()
{
    ssize_t bytesSent = 0;

    std::lock_guard<std::mutex> lock(m_writeLock);
    if (!m_writeBuffer.empty())
    {
        bytesSent = ::send(static_cast<int>(m_socket), m_writeBuffer.c_str(), static_cast<int>(m_writeBuffer.size()), 0);

        if (bytesSent > 0)
        {
            m_writeBuffer.erase(0, bytesSent);
        }
    }

    return bytesSent;
}
